from src.model.shufflenet_v2 import ShuffleNet_V2
from src.utils import Data

# 训练模型
if __name__ == '__main__':
    train_filename = r"H:\wangjianlian\data\formal_data\HER2\thrid_generation\test"
    test_filename = r"H:\wangjianlian\data\formal_data\HER2\thrid_generation\val"
    batch_size = 100
    epochs = 50
    if_save_model = True

    shuffleNet_V2 = ShuffleNet_V2()
    model = shuffleNet_V2.build(input_shape = (256, 256, 3), classes = 4)
    model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])

    data = Data()

    for epoch in range(epochs):
        print("************ 第%d代 *************"%(epoch + 1))
        all_loss = 0
        all_metrics = 0
        for index, trained_imgs, trained_label in data.read_image(train_filename, batch_size, shuffle = True):
            x_train = trained_imgs
            y_train = trained_label[0]
            loss_and_metrics  = model.train_on_batch(x_train, y_train) # 应该是返回损失值和metrics
            print("训练第%d批，loss值:%f" % (index, loss_and_metrics[0]))
            print("训练第%d批，metrics值:%f" % (index, loss_and_metrics[1]))
            all_loss += loss_and_metrics[0]
            all_metrics += loss_and_metrics[1]

        print("************ 训练：平均值 *************")
        print("训练平均loss值:%f" % (all_loss / index))
        print("训练平均metrics值:%f" % (all_metrics / index))

        all_loss = 0
        all_metrics = 0
        for index, tested_imgs, tested_label in data.read_image(test_filename, batch_size, shuffle = True):
            x_test = tested_imgs
            y_test = tested_label
            loss_and_metrics = model.evaluate(x_test, y_test, batch_size=batch_size)
            # print(loss_and_metrics)
            print("测试第%d批，loss值:%f"%(index, loss_and_metrics[0]))
            print("测试第%d批，metrics值:%f"%(index, loss_and_metrics[1]))
            all_loss += loss_and_metrics[0]
            all_metrics += loss_and_metrics[1]

        print("************ 测试：平均值 *************" )
        print("测试平均loss值:%f"%(all_loss/index))
        print("测试平均metrics值:%f"%(all_metrics/index))

        if if_save_model:
            if all_metrics/index > 0.95:
                model.save("./weights/" + str(epoch + 1) + '_ShuffleNet_V2_model.h5')